{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Finetuning the library models for question-answering on SQuAD (Bert, XLM, XLNet).\"\"\"\n",
    "\n",
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import glob\n",
    "from box import Box\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,\n",
    "                              TensorDataset)\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "from transformers import (WEIGHTS_NAME, BertConfig,\n",
    "                                  BertForQuestionAnswering, BertTokenizer,\n",
    "                                  XLMConfig, XLMForQuestionAnswering,\n",
    "                                  XLMTokenizer, XLNetConfig,\n",
    "                                  XLNetForQuestionAnswering,\n",
    "                                  XLNetTokenizer,\n",
    "                                  DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)\n",
    "\n",
    "from transformers import AdamW, WarmupLinearSchedule\n",
    "\n",
    "from utils_squad import (read_squad_examples, convert_examples_to_features,\n",
    "                         RawResult, write_predictions,\n",
    "                         RawResultExtended, write_predictions_extended)\n",
    "\n",
    "# The follwing import is the official SQuAD evaluation script (2.0).\n",
    "# You can remove it from the dependencies if you are using this script outside of the library\n",
    "# We've added it here for automated tests (see examples/test_examples.py file)\n",
    "from utils_squad_evaluate import EVAL_OPTS, main as evaluate_on_squad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) \\\n",
    "                  for conf in (BertConfig, XLNetConfig, XLMConfig)), ())\n",
    "\n",
    "MODEL_CLASSES = {\n",
    "    'bert': (BertConfig, BertForQuestionAnswering, BertTokenizer),\n",
    "    'xlnet': (XLNetConfig, XLNetForQuestionAnswering, XLNetTokenizer),\n",
    "    'xlm': (XLMConfig, XLMForQuestionAnswering, XLMTokenizer),\n",
    "    'distilbert': (DistilBertConfig, DistilBertForQuestionAnswering, DistilBertTokenizer)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_seed(args):\n",
    "    random.seed(args.seed)\n",
    "    np.random.seed(args.seed)\n",
    "    torch.manual_seed(args.seed)\n",
    "    if args.n_gpu > 0:\n",
    "        torch.cuda.manual_seed_all(args.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_list(tensor):\n",
    "    return tensor.detach().cpu().tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(args, train_dataset, model, tokenizer):\n",
    "    \"\"\" Train the model \"\"\"\n",
    "    if args.local_rank in [-1, 0]:\n",
    "        tb_writer = SummaryWriter()\n",
    "\n",
    "    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)\n",
    "    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)\n",
    "    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)\n",
    "\n",
    "    if args.max_steps > 0:\n",
    "        t_total = args.max_steps\n",
    "        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1\n",
    "    else:\n",
    "        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs\n",
    "\n",
    "    # Prepare optimizer and schedule (linear warmup and decay)\n",
    "    no_decay = ['bias', 'LayerNorm.weight']\n",
    "    optimizer_grouped_parameters = [\n",
    "        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},\n",
    "        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "        ]\n",
    "    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)\n",
    "    scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)\n",
    "    if args.fp16:\n",
    "        try:\n",
    "            from apex import amp\n",
    "        except ImportError:\n",
    "            raise ImportError(\"Please install apex from https://www.github.com/nvidia/apex to use fp16 training.\")\n",
    "        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)\n",
    "\n",
    "    # multi-gpu training (should be after apex fp16 initialization)\n",
    "    if args.n_gpu > 1:\n",
    "        model = torch.nn.DataParallel(model)\n",
    "\n",
    "    # Distributed training (should be after apex fp16 initialization)\n",
    "    if args.local_rank != -1:\n",
    "        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],\n",
    "                                                          output_device=args.local_rank,\n",
    "                                                          find_unused_parameters=True)\n",
    "\n",
    "    # Train!\n",
    "    logger.info(\"***** Running training *****\")\n",
    "    logger.info(\"  Num examples = %d\", len(train_dataset))\n",
    "    logger.info(\"  Num Epochs = %d\", args.num_train_epochs)\n",
    "    logger.info(\"  Instantaneous batch size per GPU = %d\", args.per_gpu_train_batch_size)\n",
    "    logger.info(\"  Total train batch size (w. parallel, distributed & accumulation) = %d\",\n",
    "                   args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))\n",
    "    logger.info(\"  Gradient Accumulation steps = %d\", args.gradient_accumulation_steps)\n",
    "    logger.info(\"  Total optimization steps = %d\", t_total)\n",
    "\n",
    "    global_step = 0\n",
    "    tr_loss, logging_loss = 0.0, 0.0\n",
    "    model.zero_grad()\n",
    "    train_iterator = trange(int(args.num_train_epochs), desc=\"Epoch\", disable=args.local_rank not in [-1, 0])\n",
    "    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)\n",
    "    for _ in train_iterator:\n",
    "        epoch_iterator = tqdm(train_dataloader, desc=\"Iteration\", disable=args.local_rank not in [-1, 0])\n",
    "        for step, batch in enumerate(epoch_iterator):\n",
    "            model.train()\n",
    "            batch = tuple(t.to(args.device) for t in batch)\n",
    "            inputs = {'input_ids':       batch[0],\n",
    "                      'attention_mask':  batch[1],\n",
    "                      'start_positions': batch[3], \n",
    "                      'end_positions':   batch[4]}\n",
    "            if args.model_type != 'distilbert':\n",
    "                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids\n",
    "            if args.model_type in ['xlnet', 'xlm']:\n",
    "                inputs.update({'cls_index': batch[5],\n",
    "                               'p_mask':       batch[6]})\n",
    "            outputs = model(**inputs)\n",
    "            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)\n",
    "\n",
    "            if args.n_gpu > 1:\n",
    "                loss = loss.mean() # mean() to average on multi-gpu parallel (not distributed) training\n",
    "            if args.gradient_accumulation_steps > 1:\n",
    "                loss = loss / args.gradient_accumulation_steps\n",
    "\n",
    "            if args.fp16:\n",
    "                with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "                    scaled_loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)\n",
    "            else:\n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)\n",
    "\n",
    "            tr_loss += loss.item()\n",
    "            if (step + 1) % args.gradient_accumulation_steps == 0:\n",
    "                optimizer.step()\n",
    "                scheduler.step()  # Update learning rate schedule\n",
    "                model.zero_grad()\n",
    "                global_step += 1\n",
    "\n",
    "                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:\n",
    "                    # Log metrics\n",
    "                    if args.local_rank == -1 and args.evaluate_during_training:  # Only evaluate when single GPU otherwise metrics may not average well\n",
    "                        results = evaluate(args, model, tokenizer)\n",
    "                        for key, value in results.items():\n",
    "                            tb_writer.add_scalar('eval_{}'.format(key), value, global_step)\n",
    "                    tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)\n",
    "                    tb_writer.add_scalar('loss', (tr_loss - logging_loss)/args.logging_steps, global_step)\n",
    "                    logging_loss = tr_loss\n",
    "\n",
    "                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:\n",
    "                    # Save model checkpoint\n",
    "                    output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step))\n",
    "                    if not os.path.exists(output_dir):\n",
    "                        os.makedirs(output_dir)\n",
    "                    model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training\n",
    "                    model_to_save.save_pretrained(output_dir)\n",
    "                    torch.save(args, os.path.join(output_dir, 'training_args.bin'))\n",
    "                    logger.info(\"Saving model checkpoint to %s\", output_dir)\n",
    "\n",
    "            if args.max_steps > 0 and global_step > args.max_steps:\n",
    "                epoch_iterator.close()\n",
    "                break\n",
    "        if args.max_steps > 0 and global_step > args.max_steps:\n",
    "            train_iterator.close()\n",
    "            break\n",
    "\n",
    "    if args.local_rank in [-1, 0]:\n",
    "        tb_writer.close()\n",
    "\n",
    "    return global_step, tr_loss / global_step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(args, model, tokenizer, prefix=\"\"):\n",
    "    dataset, examples, features = load_and_cache_examples(args, tokenizer, evaluate=True, output_examples=True)\n",
    "\n",
    "    if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:\n",
    "        os.makedirs(args.output_dir)\n",
    "\n",
    "    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)\n",
    "    # Note that DistributedSampler samples randomly\n",
    "    eval_sampler = SequentialSampler(dataset) if args.local_rank == -1 else DistributedSampler(dataset)\n",
    "    eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)\n",
    "\n",
    "    # Eval!\n",
    "    logger.info(\"***** Running evaluation {} *****\".format(prefix))\n",
    "    logger.info(\"  Num examples = %d\", len(dataset))\n",
    "    logger.info(\"  Batch size = %d\", args.eval_batch_size)\n",
    "    all_results = []\n",
    "    for batch in tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
    "        model.eval()\n",
    "        batch = tuple(t.to(args.device) for t in batch)\n",
    "        with torch.no_grad():\n",
    "            inputs = {'input_ids':      batch[0],\n",
    "                      'attention_mask': batch[1],\n",
    "                      }\n",
    "            if args.model_type != 'distilbert':\n",
    "                inputs['token_type_ids'] = batch[2] if args.model_type in ['bert', 'xlnet'] else None  # XLM, DistilBERT and RoBERTa don't use segment_ids\n",
    "            example_indices = batch[3]\n",
    "            if args.model_type in ['xlnet', 'xlm']:\n",
    "                inputs.update({'cls_index': batch[4],\n",
    "                               'p_mask':    batch[5]})\n",
    "            outputs = model(**inputs)\n",
    "\n",
    "        for i, example_index in enumerate(example_indices):\n",
    "            eval_feature = features[example_index.item()]\n",
    "            unique_id = int(eval_feature.unique_id)\n",
    "            if args.model_type in ['xlnet', 'xlm']:\n",
    "                # XLNet uses a more complex post-processing procedure\n",
    "                result = RawResultExtended(unique_id            = unique_id,\n",
    "                                           start_top_log_probs  = to_list(outputs[0][i]),\n",
    "                                           start_top_index      = to_list(outputs[1][i]),\n",
    "                                           end_top_log_probs    = to_list(outputs[2][i]),\n",
    "                                           end_top_index        = to_list(outputs[3][i]),\n",
    "                                           cls_logits           = to_list(outputs[4][i]))\n",
    "            else:\n",
    "                result = RawResult(unique_id    = unique_id,\n",
    "                                   start_logits = to_list(outputs[0][i]),\n",
    "                                   end_logits   = to_list(outputs[1][i]))\n",
    "            all_results.append(result)\n",
    "\n",
    "    # Compute predictions\n",
    "    output_prediction_file = os.path.join(args.output_dir, \"predictions_{}.json\".format(prefix))\n",
    "    output_nbest_file = os.path.join(args.output_dir, \"nbest_predictions_{}.json\".format(prefix))\n",
    "    if args.version_2_with_negative:\n",
    "        output_null_log_odds_file = os.path.join(args.output_dir, \"null_odds_{}.json\".format(prefix))\n",
    "    else:\n",
    "        output_null_log_odds_file = None\n",
    "\n",
    "    if args.model_type in ['xlnet', 'xlm']:\n",
    "        # XLNet uses a more complex post-processing procedure\n",
    "        write_predictions_extended(examples, features, all_results, args.n_best_size,\n",
    "                        args.max_answer_length, output_prediction_file,\n",
    "                        output_nbest_file, output_null_log_odds_file, args.predict_file,\n",
    "                        model.config.start_n_top, model.config.end_n_top,\n",
    "                        args.version_2_with_negative, tokenizer, args.verbose_logging)\n",
    "    else:\n",
    "        write_predictions(examples, features, all_results, args.n_best_size,\n",
    "                        args.max_answer_length, args.do_lower_case, output_prediction_file,\n",
    "                        output_nbest_file, output_null_log_odds_file, args.verbose_logging,\n",
    "                        args.version_2_with_negative, args.null_score_diff_threshold)\n",
    "\n",
    "    # Evaluate with the official SQuAD script\n",
    "    evaluate_options = EVAL_OPTS(data_file=args.predict_file,\n",
    "                                 pred_file=output_prediction_file,\n",
    "                                 na_prob_file=output_null_log_odds_file)\n",
    "    results = evaluate_on_squad(evaluate_options)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False):\n",
    "    if args.local_rank not in [-1, 0] and not evaluate:\n",
    "        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache\n",
    "\n",
    "    # Load data features from cache or dataset file\n",
    "    input_file = args.predict_file if evaluate else args.train_file\n",
    "    cached_features_file = os.path.join(os.path.dirname(input_file), 'cached_{}_{}_{}'.format(\n",
    "        'dev' if evaluate else 'train',\n",
    "        list(filter(None, args.model_name_or_path.split('/'))).pop(),\n",
    "        str(args.max_seq_length)))\n",
    "    if os.path.exists(cached_features_file) and not args.overwrite_cache and not output_examples:\n",
    "        logger.info(\"Loading features from cached file %s\", cached_features_file)\n",
    "        features = torch.load(cached_features_file)\n",
    "    else:\n",
    "        logger.info(\"Creating features from dataset file at %s\", input_file)\n",
    "        examples = read_squad_examples(input_file=input_file,\n",
    "                                                is_training=not evaluate,\n",
    "                                                version_2_with_negative=args.version_2_with_negative)\n",
    "        features = convert_examples_to_features(examples=examples,\n",
    "                                                tokenizer=tokenizer,\n",
    "                                                max_seq_length=args.max_seq_length,\n",
    "                                                doc_stride=args.doc_stride,\n",
    "                                                max_query_length=args.max_query_length,\n",
    "                                                is_training=not evaluate)\n",
    "        if args.local_rank in [-1, 0]:\n",
    "            logger.info(\"Saving features into cached file %s\", cached_features_file)\n",
    "            torch.save(features, cached_features_file)\n",
    "\n",
    "    if args.local_rank == 0 and not evaluate:\n",
    "        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache\n",
    "\n",
    "    # Convert to Tensors and build dataset\n",
    "    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)\n",
    "    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)\n",
    "    all_cls_index = torch.tensor([f.cls_index for f in features], dtype=torch.long)\n",
    "    all_p_mask = torch.tensor([f.p_mask for f in features], dtype=torch.float)\n",
    "    if evaluate:\n",
    "        all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)\n",
    "        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,\n",
    "                                all_example_index, all_cls_index, all_p_mask)\n",
    "    else:\n",
    "        all_start_positions = torch.tensor([f.start_position for f in features], dtype=torch.long)\n",
    "        all_end_positions = torch.tensor([f.end_position for f in features], dtype=torch.long)\n",
    "        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,\n",
    "                                all_start_positions, all_end_positions,\n",
    "                                all_cls_index, all_p_mask)\n",
    "\n",
    "    if output_examples:\n",
    "        return dataset, examples, features\n",
    "    return dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Box({\n",
    "    ## Required parameters\n",
    "    \"train_file\": \"./data/train-v1.1.json\",\n",
    "    \"predict_file\": \"./data/dev-v1.1.json\",\n",
    "    \"model_type\": 'distilbert',\n",
    "    \"model_name_or_path\": 'distilbert-base-uncased',\n",
    "    \"output_dir\": \"./output\",\n",
    "\n",
    "    ## Other parameters\n",
    "    \"config_name\": \"\", # Pretrained config name or path if not the same as model_name\n",
    "    \"tokenizer_name\": \"\", # Pretrained tokenizer name or path if not the same as model_name\n",
    "    \"cache_dir\": \"\", # Where do you want to store the pre-trained models downloaded from s3\"\n",
    "    \n",
    "    \"version_2_with_negative\": False, # If true, the SQuAD examples contain some that do not have an answer.\n",
    "    \"null_score_diff_threshold\": 0.0, # If null_score - best_non_null is greater than the threshold predict null.\n",
    "   \n",
    "    \"max_seq_length\": 384, # The maximum total input sequence length after WordPiece tokenization.\n",
    "                           # Sequences longer than this will be truncated, and sequences shorter than this will be padded.\n",
    "    \"doc_stride\": 128, # When splitting up a long document into chunks, how much stride to take between chunks.\n",
    "    \"max_query_length\": 64, # The maximum number of tokens for the question. Questions longer than this will\n",
    "                            # be truncated to this length.\n",
    "    \"do_train\": True,\n",
    "    \"do_eval\": True,\n",
    "    \"evaluate_during_training\": False,\n",
    "    \"do_lower_case\": True, \n",
    "    \n",
    "    \"per_gpu_train_batch_size\": 24,\n",
    "    \"per_gpu_eval_batch_size\": 24,\n",
    "    \"learning_rate\": 3e-5,\n",
    "    \"gradient_accumulation_steps\": 1, # Number of updates steps to accumulate before performing a backward/update pass.\n",
    "    \"weight_decay\": 0.0,\n",
    "    \"adam_epsilon\": 1e-8,\n",
    "    \"max_grad_norm\": 1.0,\n",
    "    \"num_train_epochs\": 3.0,\n",
    "    \"max_steps\": -1, # If > 0: set total number of training steps to perform. Override num_train_epochs.\")\n",
    "    \"warmup_steps\": 0,\n",
    "    \"n_best_size\": 20, # The total number of n-best predictions to generate in the nbest_predictions.json output file.\n",
    "    \"max_answer_length\": 30, # The maximum length of an answer that can be generated. This is needed because the start\n",
    "                               # and end predictions are not conditioned on one another.\n",
    "    \"verbose_logging\": False,\n",
    "    \"logging_steps\": 50,\n",
    "    \"save_steps\": -1,\n",
    "    \"eval_all_checkpoints\": False,\n",
    "    \"no_cuda\": False,\n",
    "    \"overwrite_output_dir\": True,\n",
    "    \"overwrite_cache\": True,\n",
    "    \"seed\": 42,\n",
    "\n",
    "    \"local_rank\": -1, # Local_rank for distributed training on gpus\n",
    "    \"fp16\": True, # Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit\")\n",
    "    \"fp16_opt_level\": 'O1',\n",
    "    \"server_ip\": '',\n",
    "    \"server_port\": ''\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    \n",
    "    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir:\n",
    "        raise ValueError(\"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.\".format(args.output_dir))\n",
    "\n",
    "    # Setup distant debugging if needed\n",
    "    if args.server_ip and args.server_port:\n",
    "        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script\n",
    "        import ptvsd\n",
    "        print(\"Waiting for debugger attach\")\n",
    "        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)\n",
    "        ptvsd.wait_for_attach()\n",
    "\n",
    "    # Setup CUDA, GPU & distributed training\n",
    "    if args.local_rank == -1 or args.no_cuda:\n",
    "        device = torch.device(\"cuda\" if torch.cuda.is_available() and not args.no_cuda else \"cpu\")\n",
    "        args.n_gpu = torch.cuda.device_count()\n",
    "    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs\n",
    "        torch.cuda.set_device(args.local_rank)\n",
    "        device = torch.device(\"cuda\", args.local_rank)\n",
    "        torch.distributed.init_process_group(backend='nccl')\n",
    "        args.n_gpu = 1\n",
    "    args.device = device\n",
    "\n",
    "    # Setup logging\n",
    "    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
    "                        datefmt = '%m/%d/%Y %H:%M:%S',\n",
    "                        level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN)\n",
    "    logger.warning(\"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s\",\n",
    "                    args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)\n",
    "\n",
    "    # Set seed\n",
    "    set_seed(args)\n",
    "\n",
    "    # Load pretrained model and tokenizer\n",
    "    if args.local_rank not in [-1, 0]:\n",
    "        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab\n",
    "\n",
    "    args.model_type = args.model_type.lower()\n",
    "    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]\n",
    "    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)\n",
    "    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, do_lower_case=args.do_lower_case)\n",
    "    model = model_class.from_pretrained(args.model_name_or_path, from_tf=bool('.ckpt' in args.model_name_or_path), config=config)\n",
    "\n",
    "    if args.local_rank == 0:\n",
    "        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab\n",
    "\n",
    "    model.to(args.device)\n",
    "\n",
    "    logger.info(\"Training/evaluation parameters %s\", args)\n",
    "\n",
    "    # Training\n",
    "    if args.do_train:\n",
    "        train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)\n",
    "        global_step, tr_loss = train(args, train_dataset, model, tokenizer)\n",
    "        logger.info(\" global_step = %s, average loss = %s\", global_step, tr_loss)\n",
    "\n",
    "\n",
    "    # Save the trained model and the tokenizer\n",
    "    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):\n",
    "        # Create output directory if needed\n",
    "        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:\n",
    "            os.makedirs(args.output_dir)\n",
    "\n",
    "        logger.info(\"Saving model checkpoint to %s\", args.output_dir)\n",
    "        # Save a trained model, configuration and tokenizer using `save_pretrained()`.\n",
    "        # They can then be reloaded using `from_pretrained()`\n",
    "        model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training\n",
    "        model_to_save.save_pretrained(args.output_dir)\n",
    "        tokenizer.save_pretrained(args.output_dir)\n",
    "\n",
    "        # Good practice: save your training arguments together with the trained model\n",
    "        torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))\n",
    "\n",
    "        # Load a trained model and vocabulary that you have fine-tuned\n",
    "        model = model_class.from_pretrained(args.output_dir)\n",
    "        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)\n",
    "        model.to(args.device)\n",
    "\n",
    "\n",
    "    # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory\n",
    "    results = {}\n",
    "    if args.do_eval and args.local_rank in [-1, 0]:\n",
    "        checkpoints = [args.output_dir]\n",
    "        if args.eval_all_checkpoints:\n",
    "            checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))\n",
    "            logging.getLogger(\"transformers.modeling_utils\").setLevel(logging.WARN)  # Reduce model loading logs\n",
    "\n",
    "        logger.info(\"Evaluate the following checkpoints: %s\", checkpoints)\n",
    "\n",
    "        for checkpoint in checkpoints:\n",
    "            # Reload the model\n",
    "            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else \"\"\n",
    "            model = model_class.from_pretrained(checkpoint)\n",
    "            model.to(args.device)\n",
    "\n",
    "            # Evaluate\n",
    "            result = evaluate(args, model, tokenizer, prefix=global_step)\n",
    "\n",
    "            result = dict((k + ('_{}'.format(global_step) if global_step else ''), v) for k, v in result.items())\n",
    "            results.update(result)\n",
    "\n",
    "    logger.info(\"Results: {}\".format(results))\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "     main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
